"""
强制工具执行指定函数
"""
from pprint import pprint
from langchain_core.tools import tool

from models import get_ds_model_client


@tool
def add(a: int, b: int):
    """add two numbers

    Args:
        a: first number,
        b: second number
    """
    return a + b


@tool
def multiply(a: int, b: int):
    """multiply two numbers

    Args:
        a: first number,
        b: second number
    """
    return a * b


tools = [add, multiply]
llm = get_ds_model_client()
# 还可以通过将“any”（或“required”，这是OpenAI特定的）关键字传递给tool_choice参数，强制我们的工具选择至少一个工具。参数parallel_tool_calls, 禁用并行多次执行工具
llm_bind_tools = llm.bind_tools(tools, tool_choice="multiply")
# llm_bind_tools = llm.bind_tools(tools, tool_choice="any", parallel_tool_calls=False)

query = "what is 2 + 4?"
pprint(llm_bind_tools.invoke(query))
